from typing import Optional
from asyncio import Queue
import os
from urllib.parse import urlparse

# 恶意域名数据库文件路径
MALICIOUS_DOMAINS_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))), 'docs', 'malicious_provenance.txt')

# 全局变量，用于存储加载的恶意域名集合
MALICIOUS_DOMAINS = set()

# 在模块加载时读取恶意域名数据库
try:
    with open(MALICIOUS_DOMAINS_PATH, 'r', encoding='utf-8') as f:
        for line in f:
            domain = line.strip().lower() # 读取并转换为小写
            if domain:
                MALICIOUS_DOMAINS.add(domain)
    print(f"[Malicious Referrer Detector] 成功加载 {len(MALICIOUS_DOMAINS)} 个恶意域名。")
except FileNotFoundError:
    print(f"[Malicious Referrer Detector] 警告: 未找到恶意域名数据库文件: {MALICIOUS_DOMAINS_PATH}。恶意来源检测功能可能受限。")
except Exception as e:
    print(f"[Malicious Referrer Detector] 警告: 加载恶意域名数据库文件时发生错误: {e}。恶意来源检测功能可能受限。")

async def _send_status(status_queue: Optional[Queue], tool_name: str, message_type: str, data: dict):
    """辅助函数：向状态队列发送消息。"""
    if status_queue:
        await status_queue.put({
            "type": message_type,
            "step": tool_name,
            "data": data
        })

async def is_malicious_referrer(referrer: Optional[str], status_queue: Optional[Queue] = None) -> dict:
    """
    判断referrer是否来自恶意来源。
    """
    tool_name = "Malicious_Referrer_Detector"
    await _send_status(status_queue, tool_name, "tool_start", {"referrer": referrer})

    is_malicious = False
    reason = "Referrer未在已知恶意来源数据库中发现。"
    extracted_domain = None

    if not referrer:
        reason = "Referrer字符串为空或None。"
    else:
        try:
            # 尝试解析URL，提取域名
            parsed_url = urlparse(referrer)
            extracted_domain = parsed_url.netloc
            
            # 如果netloc为空，尝试直接使用referrer作为域名（例如，只有域名没有协议）
            if not extracted_domain:
                extracted_domain = referrer.split('/')[0].split(':')[0].lower()
            else:
                extracted_domain = extracted_domain.lower()

            if extracted_domain in MALICIOUS_DOMAINS:
                is_malicious = True
                reason = f"Referrer域名 '{extracted_domain}' 在已知恶意来源数据库中发现。"
            
        except Exception as e:
            reason = f"解析Referrer域名时发生错误: {e}"

    analysis_result = {
        "is_malicious_referrer": is_malicious,
        "reason": reason,
        "extracted_domain": extracted_domain
    }

    await _send_status(status_queue, tool_name, "tool_end", {"referrer": referrer, "result": analysis_result})
    return {"source": "Internal-Malicious-Referrer-Detector", "data": analysis_result, "status": "success"}
